import os
import csv
# Replace with the path to your directory
path = "/home/eric/dev/dsp/data/abelian"

# all
with open('output.csv', 'w') as output:
    writer = csv.writer(output)
    writer.writerow(["group", "v", "k1", "k2", "e",
                    "lambda", "k1set", "k2set", "eset"])
    for filename in os.listdir(path):
        if filename.endswith(".txt"):
            print("processing", os.path.join(path, filename))
            # Extract the order and index from the filename
            parts = filename.split("-")
            order = int(parts[0])
            # if order % 2 != 0:
            #     continue
            # Remove the ".txt" extension from the second part
            index = int(parts[1][:-4])

            # Read in the pairs from the file
            with open(os.path.join(path, filename), "r") as f:
                f.readline()
                second_line = f.readline().strip()
                count = int(second_line.split(": ")[1])

                for i in range(3):  # Skip the next 3 lines
                    f.readline()

                pairs = []
                current_pair = ""
                for i in range(count):  # Read in the pairs
                    while True:  # Concatenate lines until a complete pair is read
                        line = f.readline().strip()
                        current_pair += line
                        if line.endswith("]") or line.endswith("] ],"):
                            break

                    if current_pair.endswith(","):
                        current_pair = current_pair[:-1]
                    # Convert the string to a pair object
                    pair = eval(current_pair)
                    pairs.append(pair)
                    current_pair = ""

                # Do something with the order, index, and pairs here
                print("File {} contains order {}, index {}, count {}, and pairs {}".format(
                    filename, order, index, count, pairs))

                # save to output
                for pair in pairs:
                    k1set = set(pair[0])
                    k2set = set(pair[1])
                    eset = k1set.intersection(k2set)
                    k1 = len(k1set)
                    k2 = len(k2set)
                    e = len(eset)
                    l = int((k1*k2-e)/(order-1))
                    writer.writerow(
                        [f"{order}-{index}", order, k1, k2, e, l, k1set, k2set, eset])
# Z2n
with open('output-Z2n.csv', 'w') as output:
    writer = csv.writer(output)
    writer.writerow(["group", "v", "k1", "k2", "e",
                    "lambda", "k1set", "k2set", "eset"])
    for filename in os.listdir(path):
        if filename.endswith(".txt"):
            print("processing", os.path.join(path, filename))
            # Extract the order and index from the filename
            parts = filename.split("-")
            order = int(parts[0])
            if order % 2 != 0:
                continue
            # Remove the ".txt" extension from the second part
            index = int(parts[1][:-4])

            # Read in the pairs from the file
            with open(os.path.join(path, filename), "r") as f:
                f.readline()
                second_line = f.readline().strip()
                count = int(second_line.split(": ")[1])

                for i in range(3):  # Skip the next 3 lines
                    f.readline()

                pairs = []
                current_pair = ""
                for i in range(count):  # Read in the pairs
                    while True:  # Concatenate lines until a complete pair is read
                        line = f.readline().strip()
                        current_pair += line
                        if line.endswith("]") or line.endswith("] ],"):
                            break

                    if current_pair.endswith(","):
                        current_pair = current_pair[:-1]
                    # Convert the string to a pair object
                    pair = eval(current_pair)
                    pairs.append(pair)
                    current_pair = ""

                # Do something with the order, index, and pairs here
                print("File {} contains order {}, index {}, count {}, and pairs {}".format(
                    filename, order, index, count, pairs))

                # save to output
                for pair in pairs:
                    k1set = set(pair[0])
                    k2set = set(pair[1])
                    eset = k1set.intersection(k2set)
                    k1 = len(k1set)
                    k2 = len(k2set)
                    e = len(eset)
                    l = int((k1*k2-e)/(order-1))
                    writer.writerow(
                        [f"{order}-{index}", order, k1, k2, e, l, k1set, k2set, eset])

# Z3n
with open('output-Z3n.csv', 'w') as output:
    writer = csv.writer(output)
    writer.writerow(["group", "v", "k1", "k2", "e",
                    "lambda", "k1set", "k2set", "eset"])
    for filename in os.listdir(path):
        if filename.endswith(".txt"):
            print("processing", os.path.join(path, filename))
            # Extract the order and index from the filename
            parts = filename.split("-")
            order = int(parts[0])
            if order % 3 != 0:
                continue
            # Remove the ".txt" extension from the second part
            index = int(parts[1][:-4])

            # Read in the pairs from the file
            with open(os.path.join(path, filename), "r") as f:
                f.readline()
                second_line = f.readline().strip()
                count = int(second_line.split(": ")[1])

                for i in range(3):  # Skip the next 3 lines
                    f.readline()

                pairs = []
                current_pair = ""
                for i in range(count):  # Read in the pairs
                    while True:  # Concatenate lines until a complete pair is read
                        line = f.readline().strip()
                        current_pair += line
                        if line.endswith("]") or line.endswith("] ],"):
                            break

                    if current_pair.endswith(","):
                        current_pair = current_pair[:-1]
                    # Convert the string to a pair object
                    pair = eval(current_pair)
                    pairs.append(pair)
                    current_pair = ""

                # Do something with the order, index, and pairs here
                print("File {} contains order {}, index {}, count {}, and pairs {}".format(
                    filename, order, index, count, pairs))

                # save to output
                for pair in pairs:
                    k1set = set(pair[0])
                    k2set = set(pair[1])
                    eset = k1set.intersection(k2set)
                    k1 = len(k1set)
                    k2 = len(k2set)
                    e = len(eset)
                    l = int((k1*k2-e)/(order-1))
                    writer.writerow(
                        [f"{order}-{index}", order, k1, k2, e, l, k1set, k2set, eset])

def is_prime(n):
    if n <= 1:
        return False
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            return False
    return True


# Zp
with open('output-Zp.csv', 'w') as output:
    writer = csv.writer(output)
    writer.writerow(["group", "v", "k1", "k2", "e",
                    "lambda", "k1set", "k2set", "eset"])
    for filename in os.listdir(path):
        if filename.endswith(".txt"):
            print("processing", os.path.join(path, filename))
            # Extract the order and index from the filename
            parts = filename.split("-")
            order = int(parts[0])
            if not is_prime(order):
                continue
            # Remove the ".txt" extension from the second part
            index = int(parts[1][:-4])

            # Read in the pairs from the file
            with open(os.path.join(path, filename), "r") as f:
                f.readline()
                second_line = f.readline().strip()
                count = int(second_line.split(": ")[1])

                for i in range(3):  # Skip the next 3 lines
                    f.readline()

                pairs = []
                current_pair = ""
                for i in range(count):  # Read in the pairs
                    while True:  # Concatenate lines until a complete pair is read
                        line = f.readline().strip()
                        current_pair += line
                        if line.endswith("]") or line.endswith("] ],"):
                            break

                    if current_pair.endswith(","):
                        current_pair = current_pair[:-1]
                    # Convert the string to a pair object
                    pair = eval(current_pair)
                    pairs.append(pair)
                    current_pair = ""

                # Do something with the order, index, and pairs here
                print("File {} contains order {}, index {}, count {}, and pairs {}".format(
                    filename, order, index, count, pairs))

                # save to output
                for pair in pairs:
                    k1set = set(pair[0])
                    k2set = set(pair[1])
                    eset = k1set.intersection(k2set)
                    k1 = len(k1set)
                    k2 = len(k2set)
                    e = len(eset)
                    l = int((k1*k2-e)/(order-1))
                    writer.writerow(
                        [f"{order}-{index}", order, k1, k2, e, l, k1set, k2set, eset])
